// 文件路径: dbms/service/AIService.java
package com.dbms.service;

import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import okhttp3.RequestBody;
import okhttp3.Response;
import org.json.JSONArray;
import org.json.JSONObject;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.Proxy;
import java.util.concurrent.TimeUnit;

/**
 * AI服务层核心类：负责与Google Gemini AI模型交互，根据用户自然语言描述生成符合规范的UDF（用户自定义函数）Java代码。
 * 【关键特性】
 * 1. 支持HTTP代理配置：解决需要代理环境（如国内访问Google服务）的网络连接问题；
 * 2. 严格的请求/响应处理：包含API密钥校验、超时控制、错误捕获与友好提示；
 * 3. 规范的代码生成：强制UDF代码符合DBMS系统的包结构、类名、接口实现要求。
 */
public class AIService {

    /** Gemini AI API请求地址：指定使用1.5 Flash最新版本模型（兼顾速度与生成质量） */
    private static final String API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-latest:generateContent";
    /** Gemini API密钥：从系统环境变量获取，避免硬编码泄露风险 */
    private final String apiKey;
    /** OKHttp客户端：用于发送HTTP请求，已配置代理和超时策略 */
    private final OkHttpClient client;

    /**
     * 构造方法：初始化API密钥和OKHttp客户端（核心是代理配置）
     * 代理默认适配常见工具（Clash for Windows默认7890端口，v2rayN默认10809端口，需根据实际环境调整）
     */
    public AIService() {
        // 从系统环境变量读取Gemini API密钥（推荐部署时通过环境变量注入，安全无硬编码）
        this.apiKey = System.getenv("GEMINI_API_KEY");
        // 校验API密钥是否存在，不存在则打印警告（AI功能不可用）
        if (this.apiKey == null || this.apiKey.isEmpty()) {
            System.err.println("警告: 未找到环境变量 'GEMINI_API_KEY'。AI功能将不可用。");
        }

        // --- 【核心修改】配置HTTP代理：解决跨地区API访问问题 ---
        // Proxy.Type.HTTP：指定使用HTTP代理类型；
        // InetSocketAddress：代理服务器地址（127.0.0.1为本地代理，端口7890为Clash默认端口）
        Proxy proxy = new Proxy(Proxy.Type.HTTP, new InetSocketAddress("127.0.0.1", 7890));

        // 初始化OKHttp客户端，配置关键参数
        this.client = new OkHttpClient.Builder()
                .proxy(proxy) // 应用代理配置，无代理环境可注释此句
                .connectTimeout(60, TimeUnit.SECONDS) // 连接超时：60秒（应对网络波动）
                .readTimeout(60, TimeUnit.SECONDS) // 读取超时：60秒（等待AI生成响应）
                .writeTimeout(60, TimeUnit.SECONDS) // 写入超时：60秒（避免请求发送中断）
                .build();
    }

    /**
     * 核心方法：根据用户自然语言描述，调用Gemini AI生成UDF代码
     * @param description 用户对UDF功能的自然语言描述（如“创建字符串MD5哈希函数”）
     * @return 生成的UDF代码（成功）或错误提示信息（失败）
     */
    public String generateUdfCode(String description) {
        // 1. 先校验API密钥是否配置，未配置直接返回错误
        if (apiKey == null) {
            return "错误: Gemini API 密钥未配置。请设置 'GEMINI_API_KEY' 环境变量。";
        }

        try {
            // 2. 构建AI请求的Prompt（包含用户需求和严格的代码生成规则）
            String prompt = buildPrompt(description);

            // 3. 构造Gemini API要求的JSON请求体（遵循Google Generative AI API格式）
            // 第一层：parts数组，存放单个text类型的请求内容
            JSONObject textPart = new JSONObject().put("text", prompt);
            JSONArray parts = new JSONArray().put(textPart);
            // 第二层：content对象，包含parts数组
            JSONObject content = new JSONObject().put("parts", parts);
            // 最外层：contents数组，存放单个content对象（Gemini支持多轮对话，此处为单轮）
            JSONObject bodyJson = new JSONObject().put("contents", new JSONArray().put(content));

            // 4. 构建HTTP请求体：指定JSON格式，字符编码UTF-8
            RequestBody body = RequestBody.create(
                    bodyJson.toString(),
                    MediaType.get("application/json; charset=utf-8")
            );

            // 5. 构建完整HTTP请求：URL拼接API密钥，POST方法提交请求体
            okhttp3.Request request = new okhttp3.Request.Builder()
                    .url(API_URL + "?key=" + apiKey) // URL参数传递API密钥（Gemini API要求）
                    .post(body) // POST方法（提交生成请求）
                    .build();

            // 6. 发送请求并处理响应（try-with-resources自动关闭Response，避免资源泄露）
            try (Response response = client.newCall(request).execute()) {
                // 获取响应体字符串（无响应体时设为“无响应体”）
                String responseBody = response.body() != null ? response.body().string() : "无响应体";

                // 7. 校验HTTP响应状态：非2xx（成功）状态码返回错误信息
                if (!response.isSuccessful()) {
                    return "错误: API请求失败: " + response.code() + " " + response.message() + "\n响应: " + responseBody;
                }

                // 8. 解析AI响应，提取并返回符合要求的UDF代码
                return parseResponseAndExtractCode(responseBody);
            }

        } catch (Exception e) {
            // 捕获所有异常（网络异常、JSON解析异常等），打印堆栈并返回友好提示
            e.printStackTrace();
            return "错误: 调用AI服务时发生网络异常: " + e.getMessage();
        }
    }

    /**
     * 辅助方法：解析Gemini AI的JSON响应，提取其中的Java代码块
     * 处理逻辑：优先提取```java标记包裹的代码，无标记时提取纯文本，解析失败返回原始响应
     * @param responseBody AI返回的JSON格式响应字符串
     * @return 提取的UDF代码或错误提示
     */
    private String parseResponseAndExtractCode(String responseBody) {
        try {
            // 1. 将响应字符串转为JSON对象
            JSONObject jsonResponse = new JSONObject(responseBody);

            // 2. 检查AI是否返回错误（如API密钥无效、请求超限）
            if (jsonResponse.has("error")) {
                String errorMsg = jsonResponse.getJSONObject("error").getString("message");
                return "错误: Gemini API返回错误: " + errorMsg;
            }

            // 3. 按Gemini API响应结构逐层提取生成的文本（核心路径）
            // 响应结构：candidates -> 0 -> content -> parts -> 0 -> text
            String fullText = jsonResponse.getJSONArray("candidates")
                    .getJSONObject(0) // 取第一个候选结果（AI返回的最优结果）
                    .getJSONObject("content") // 取content对象
                    .getJSONArray("parts") // 取parts数组
                    .getJSONObject(0) // 取第一个part（文本内容）
                    .getString("text"); // 取最终生成的文本

            // 4. 提取```java包裹的代码块（Markdown代码块格式）
            int start = fullText.indexOf("```java");
            if (start == -1) {
                // 兼容无java标记的情况（如仅```包裹）
                start = fullText.indexOf("```");
                if (start == -1) {
                    // 无任何代码块标记，直接返回全文
                    return fullText;
                }
                start += 3; // 跳过"```"（长度3）
            } else {
                start += 7; // 跳过"```java"（长度7）
            }

            // 5. 找到代码块结束标记```
            int end = fullText.lastIndexOf("```");
            if (end == -1 || end <= start) {
                // 无结束标记或结束标记在开始标记前，返回从开始标记后的所有内容
                return fullText.substring(start);
            }

            // 6. 提取代码块并去除前后空白（避免换行、空格问题）
            return fullText.substring(start, end).trim();

        } catch (Exception e) {
            // 解析异常（如JSON格式错误、字段缺失），返回原始响应供调试
            return "错误: 解析AI响应失败。原始响应: " + responseBody;
        }
    }

    /**
     * 辅助方法：构建AI请求的Prompt（提示词工程核心）
     * 作用：明确告诉AI生成UDF的规则，确保输出符合DBMS系统的集成要求
     * @param userDescription 用户输入的自然语言需求
     * @return 结构化、带规则约束的Prompt字符串
     */
    private String buildPrompt(String userDescription) {
        return "请根据用户的需求，生成一段可以直接在我的DBMS系统中使用的Java UDF（用户自定义函数）代码。\n" +
                "用户需求如下：\"" + userDescription + "\"\n\n" +
                "你生成的代码必须严格遵守以下规则，任何一条不满足都会导致代码无法使用：\n" +
                "1. 必须包含完整的包声明：`package com.dbms.udf.temp;`（固定包路径，不可修改）\n" +
                "2. 必须导入 `com.dbms.execution.udf.Udf` 接口（UDF需实现此接口）\n" +
                "3. 公共类名必须**严格**命名为 `CustomUdf`（系统通过类名加载UDF，不可修改）\n" +
                "4. `CustomUdf` 类必须实现 `Udf` 接口（强制接口实现，确保统一调用规范）\n" +
                "5. 必须实现 `public Object execute(Object... args)` 方法：\n" +
                "   - 所有函数核心逻辑写在此方法内；\n" +
                "   - 必须对`args`的长度（参数数量）和元素类型做严格校验；\n" +
                "   - 参数不符合预期时，抛出 `IllegalArgumentException` 并附带清晰错误信息（如“需输入1个字符串参数”）\n" +
                "6. 输出格式要求：\n" +
                "   - 只包含Java代码本身，无任何额外解释、标题、注释；\n" +
                "   - 代码必须包裹在 ```java 和 ``` 之间（Markdown代码块格式）\n\n" +
                "以下是符合要求的优秀示例（字符串反转UDF），请参考此格式生成：\n" +
                "```java\n" +
                "package com.dbms.udf.temp;\n\n" +
                "import com.dbms.execution.udf.Udf;\n\n" +
                "public class CustomUdf implements Udf {\n\n" +
                "    @Override\n" +
                "    public Object execute(Object... args) throws Exception {\n" +
                "        if (args.length != 1 || !(args[0] instanceof String)) {\n" +
                "            throw new IllegalArgumentException(\"请输入一个字符串参数\");\n" +
                "        }\n" +
                "        String input = (String) args[0];\n" +
                "        return new StringBuilder(input).reverse().toString();\n" +
                "    }\n" +
                "}\n" +
                "```";
    }
}